import json
from langchain_openai import ChatOpenAI
from langchain_community.chat_models import ChatSparkLLM
from langchain_core.output_parsers import StrOutputParser

with open('model-url.json', 'r') as file:
    data = json.load(file)


class ChatModel:
    def __init__(self, model, version, temp):
        self.model = model
        self.version = version
        self.temp = temp

    def model_choose(self):
        model = self.model
        #print(model,self.version,self.temp)
        if data[model]["class_name"] == 'ChatSparkLLM':
            return ChatSparkLLM(
                spark_api_url=data[model][self.version],
                spark_app_id=data[model]["appid"],
                spark_api_key=data[model]["ak"],
                spark_api_secret=data[model]["sk"],
                spark_llm_domain=self.version,
                temperature=self.temp,
                streaming=False,
            )
        elif data[model]["class_name"] == 'wenxin':
            pass
        else:
            return ChatOpenAI(
                model=self.version,
                temperature=self.temp,
                base_url=data[model][self.version],
                openai_api_key=data[model]["openai_api_key"]
            )

